/*
 * 3D Helmholtz Finite Element Method Solver
 * Solves: (I - κ∇²)u = f in 3D domain using linear tetrahedral elements
 * Weak form: ∫(uv + κ∇u·∇v)dΩ = ∫fv dΩ for all test functions v
 * 
 * The matrix is assembled to be symmetric positive definite with diagonal elements near 1
 * Uses structured tetrahedral mesh for simplicity
 */

#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <string.h>

#ifndef M_PI
#define M_PI 3.14159265358979323846
#endif

typedef struct {
    double x, y, z;
} Point3D;

typedef struct {
    int nodes[4];  // Node indices for tetrahedron
} Tetrahedron;

typedef struct {
    int nx, ny, nz;           // Grid dimensions
    double lx, ly, lz;        // Domain size
    double kappa;             // Helmholtz parameter
    int num_nodes;            // Total number of nodes
    int num_elements;         // Total number of elements
    Point3D* nodes;           // Node coordinates
    Tetrahedron* elements;    // Element connectivity
    double* matrix;           // Global stiffness matrix (sparse storage)
    int* ia, *ja;            // Sparse matrix structure (CSR format)
    double* rhs;             // Right-hand side vector
    double* solution;        // Solution vector
    int* boundary_nodes;     // Boundary node markers
} HelmholtzFEM3D;

// Convert 3D grid indices to node index
static inline int grid_to_node(int i, int j, int k, int nx, int ny, int nz) {
    return i + nx * (j + ny * k);
}

// Create structured tetrahedral mesh
void create_structured_tet_mesh(HelmholtzFEM3D* fem) {
    int nx = fem->nx, ny = fem->ny, nz = fem->nz;
    
    // Create nodes
    fem->num_nodes = nx * ny * nz;
    fem->nodes = malloc(fem->num_nodes * sizeof(Point3D));
    
    double dx = fem->lx / (nx - 1);
    double dy = fem->ly / (ny - 1);
    double dz = fem->lz / (nz - 1);
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = grid_to_node(i, j, k, nx, ny, nz);
                fem->nodes[idx].x = i * dx;
                fem->nodes[idx].y = j * dy;
                fem->nodes[idx].z = k * dz;
            }
        }
    }
    
    // Create elements (6 tetrahedra per cube)
    fem->num_elements = 6 * (nx - 1) * (ny - 1) * (nz - 1);
    fem->elements = malloc(fem->num_elements * sizeof(Tetrahedron));
    
    int elem_idx = 0;
    for (int k = 0; k < nz - 1; k++) {
        for (int j = 0; j < ny - 1; j++) {
            for (int i = 0; i < nx - 1; i++) {
                // Get the 8 vertices of the cube
                int n000 = grid_to_node(i,   j,   k,   nx, ny, nz);
                int n100 = grid_to_node(i+1, j,   k,   nx, ny, nz);
                int n010 = grid_to_node(i,   j+1, k,   nx, ny, nz);
                int n110 = grid_to_node(i+1, j+1, k,   nx, ny, nz);
                int n001 = grid_to_node(i,   j,   k+1, nx, ny, nz);
                int n101 = grid_to_node(i+1, j,   k+1, nx, ny, nz);
                int n011 = grid_to_node(i,   j+1, k+1, nx, ny, nz);
                int n111 = grid_to_node(i+1, j+1, k+1, nx, ny, nz);
                
                // Decompose cube into 6 tetrahedra
                // Tet 1: (0,0,0)-(1,0,0)-(0,1,0)-(0,0,1)
                fem->elements[elem_idx].nodes[0] = n000;
                fem->elements[elem_idx].nodes[1] = n100;
                fem->elements[elem_idx].nodes[2] = n010;
                fem->elements[elem_idx].nodes[3] = n001;
                elem_idx++;
                
                // Tet 2: (1,0,0)-(1,1,0)-(0,1,0)-(1,0,1)
                fem->elements[elem_idx].nodes[0] = n100;
                fem->elements[elem_idx].nodes[1] = n110;
                fem->elements[elem_idx].nodes[2] = n010;
                fem->elements[elem_idx].nodes[3] = n101;
                elem_idx++;
                
                // Tet 3: (0,1,0)-(1,1,0)-(0,1,1)-(1,0,1)
                fem->elements[elem_idx].nodes[0] = n010;
                fem->elements[elem_idx].nodes[1] = n110;
                fem->elements[elem_idx].nodes[2] = n011;
                fem->elements[elem_idx].nodes[3] = n101;
                elem_idx++;
                
                // Tet 4: (1,1,0)-(1,1,1)-(0,1,1)-(1,0,1)
                fem->elements[elem_idx].nodes[0] = n110;
                fem->elements[elem_idx].nodes[1] = n111;
                fem->elements[elem_idx].nodes[2] = n011;
                fem->elements[elem_idx].nodes[3] = n101;
                elem_idx++;
                
                // Tet 5: (0,0,1)-(1,0,1)-(0,1,1)-(0,0,0)
                fem->elements[elem_idx].nodes[0] = n001;
                fem->elements[elem_idx].nodes[1] = n101;
                fem->elements[elem_idx].nodes[2] = n011;
                fem->elements[elem_idx].nodes[3] = n000;
                elem_idx++;
                
                // Tet 6: (1,0,1)-(0,1,1)-(0,1,0)-(0,0,0)
                fem->elements[elem_idx].nodes[0] = n101;
                fem->elements[elem_idx].nodes[1] = n011;
                fem->elements[elem_idx].nodes[2] = n010;
                fem->elements[elem_idx].nodes[3] = n000;
                elem_idx++;
            }
        }
    }
    
    printf("Created mesh: %d nodes, %d elements\n", fem->num_nodes, fem->num_elements);
}

// Compute tetrahedron volume and shape function gradients
double compute_tet_volume_and_gradients(Point3D* nodes, double grad[][3]) {
    // Compute Jacobian matrix
    double J[3][3];
    for (int i = 0; i < 3; i++) {
        J[i][0] = nodes[i+1].x - nodes[0].x;
        J[i][1] = nodes[i+1].y - nodes[0].y;
        J[i][2] = nodes[i+1].z - nodes[0].z;
    }
    
    // Compute determinant
    double det = J[0][0] * (J[1][1] * J[2][2] - J[1][2] * J[2][1]) -
                 J[0][1] * (J[1][0] * J[2][2] - J[1][2] * J[2][0]) +
                 J[0][2] * (J[1][0] * J[2][1] - J[1][1] * J[2][0]);
    
    double volume = fabs(det) / 6.0;
    
    // Compute inverse Jacobian
    double invJ[3][3];
    double inv_det = 1.0 / det;
    
    invJ[0][0] = (J[1][1] * J[2][2] - J[1][2] * J[2][1]) * inv_det;
    invJ[0][1] = (J[0][2] * J[2][1] - J[0][1] * J[2][2]) * inv_det;
    invJ[0][2] = (J[0][1] * J[1][2] - J[0][2] * J[1][1]) * inv_det;
    invJ[1][0] = (J[1][2] * J[2][0] - J[1][0] * J[2][2]) * inv_det;
    invJ[1][1] = (J[0][0] * J[2][2] - J[0][2] * J[2][0]) * inv_det;
    invJ[1][2] = (J[0][2] * J[1][0] - J[0][0] * J[1][2]) * inv_det;
    invJ[2][0] = (J[1][0] * J[2][1] - J[1][1] * J[2][0]) * inv_det;
    invJ[2][1] = (J[0][1] * J[2][0] - J[0][0] * J[2][1]) * inv_det;
    invJ[2][2] = (J[0][0] * J[1][1] - J[0][1] * J[1][0]) * inv_det;
    
    // Shape function gradients in reference coordinates
    double ref_grad[4][3] = {{-1, -1, -1}, {1, 0, 0}, {0, 1, 0}, {0, 0, 1}};
    
    // Transform to physical coordinates
    for (int i = 0; i < 4; i++) {
        for (int j = 0; j < 3; j++) {
            grad[i][j] = 0.0;
            for (int k = 0; k < 3; k++) {
                grad[i][j] += invJ[j][k] * ref_grad[i][k];
            }
        }
    }
    
    return volume;
}

// Initialize sparse matrix structure (CSR format)
void init_sparse_matrix(HelmholtzFEM3D* fem) {
    int n = fem->num_nodes;
    
    // Count non-zeros per row (estimate)
    int* nnz_per_row = calloc(n, sizeof(int));
    
    // For each element, mark connections
    for (int e = 0; e < fem->num_elements; e++) {
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                int node_i = fem->elements[e].nodes[i];
                nnz_per_row[node_i]++;
            }
        }
    }
    
    // Create CSR structure
    fem->ia = malloc((n + 1) * sizeof(int));
    fem->ia[0] = 0;
    for (int i = 0; i < n; i++) {
        fem->ia[i + 1] = fem->ia[i] + nnz_per_row[i];
    }
    
    int total_nnz = fem->ia[n];
    fem->ja = malloc(total_nnz * sizeof(int));
    fem->matrix = calloc(total_nnz, sizeof(double));
    
    // Fill column indices
    int* current_pos = calloc(n, sizeof(int));
    for (int e = 0; e < fem->num_elements; e++) {
        for (int i = 0; i < 4; i++) {
            int node_i = fem->elements[e].nodes[i];
            for (int j = 0; j < 4; j++) {
                int node_j = fem->elements[e].nodes[j];
                int pos = fem->ia[node_i] + current_pos[node_i];
                fem->ja[pos] = node_j;
                current_pos[node_i]++;
            }
        }
    }
    
    printf("Sparse matrix: %d x %d with %d non-zeros\n", n, n, total_nnz);
    
    free(nnz_per_row);
    free(current_pos);
}

// Assemble global stiffness matrix and RHS
void assemble_system(HelmholtzFEM3D* fem) {
    double kappa = fem->kappa;
    
    // Clear matrix and RHS
    int total_nnz = fem->ia[fem->num_nodes];
    memset(fem->matrix, 0, total_nnz * sizeof(double));
    memset(fem->rhs, 0, fem->num_nodes * sizeof(double));
    
    // Element matrices
    double K_elem[4][4];  // Element stiffness matrix
    double M_elem[4][4];  // Element mass matrix
    double f_elem[4];     // Element RHS vector
    
    for (int e = 0; e < fem->num_elements; e++) {
        // Get element nodes
        Point3D elem_nodes[4];
        for (int i = 0; i < 4; i++) {
            elem_nodes[i] = fem->nodes[fem->elements[e].nodes[i]];
        }
        
        // Compute gradients and volume
        double grad[4][3];
        double volume = compute_tet_volume_and_gradients(elem_nodes, grad);
        
        // Compute element matrices
        for (int i = 0; i < 4; i++) {
            for (int j = 0; j < 4; j++) {
                // Stiffness matrix: ∫∇Ni·∇Nj dΩ
                K_elem[i][j] = 0.0;
                for (int d = 0; d < 3; d++) {
                    K_elem[i][j] += grad[i][d] * grad[j][d];
                }
                K_elem[i][j] *= volume;
                
                // Mass matrix: ∫NiNj dΩ
                if (i == j) {
                    M_elem[i][j] = volume / 10.0;  // Diagonal: V/10
                } else {
                    M_elem[i][j] = volume / 20.0;  // Off-diagonal: V/20
                }
            }
        }
        
        // Compute element RHS (assuming f = constant per element)
        Point3D centroid = {0, 0, 0};
        for (int i = 0; i < 4; i++) {
            centroid.x += elem_nodes[i].x / 4.0;
            centroid.y += elem_nodes[i].y / 4.0;
            centroid.z += elem_nodes[i].z / 4.0;
        }
        
        // Test function: f such that exact solution is sin(πx)*sin(πy)*sin(πz)
        double lambda = M_PI * M_PI * (1.0/(fem->lx*fem->lx) + 1.0/(fem->ly*fem->ly) + 1.0/(fem->lz*fem->lz));
        double u_exact = sin(M_PI * centroid.x / fem->lx) * 
                         sin(M_PI * centroid.y / fem->ly) * 
                         sin(M_PI * centroid.z / fem->lz);
        double f_val = (1.0 + kappa * lambda) * u_exact;
        
        for (int i = 0; i < 4; i++) {
            f_elem[i] = f_val * volume / 4.0;  // ∫f*Ni dΩ ≈ f*V/4
        }
        
        // Assemble into global system: (M + κK)u = f
        for (int i = 0; i < 4; i++) {
            int global_i = fem->elements[e].nodes[i];
            
            // Add to RHS
            fem->rhs[global_i] += f_elem[i];
            
            for (int j = 0; j < 4; j++) {
                int global_j = fem->elements[e].nodes[j];
                double coeff = M_elem[i][j] + kappa * K_elem[i][j];
                
                // Find position in sparse matrix
                for (int pos = fem->ia[global_i]; pos < fem->ia[global_i + 1]; pos++) {
                    if (fem->ja[pos] == global_j) {
                        fem->matrix[pos] += coeff;
                        break;
                    }
                }
            }
        }
    }
    
    printf("System assembled. Matrix norm estimate: %.6e\n", 
           fem->matrix[fem->ia[fem->num_nodes/2]]);
}

// Apply boundary conditions (homogeneous Dirichlet)
void apply_boundary_conditions(HelmholtzFEM3D* fem) {
    int nx = fem->nx, ny = fem->ny, nz = fem->nz;
    
    // Mark boundary nodes
    fem->boundary_nodes = calloc(fem->num_nodes, sizeof(int));
    
    for (int k = 0; k < nz; k++) {
        for (int j = 0; j < ny; j++) {
            for (int i = 0; i < nx; i++) {
                int idx = grid_to_node(i, j, k, nx, ny, nz);
                
                // Boundary if on any face
                if (i == 0 || i == nx-1 || j == 0 || j == ny-1 || k == 0 || k == nz-1) {
                    fem->boundary_nodes[idx] = 1;
                }
            }
        }
    }
    
    // Apply boundary conditions: set diagonal to 1, off-diagonals to 0, RHS to 0
    for (int i = 0; i < fem->num_nodes; i++) {
        if (fem->boundary_nodes[i]) {
            fem->rhs[i] = 0.0;
            
            // Clear row and set diagonal to 1
            for (int pos = fem->ia[i]; pos < fem->ia[i + 1]; pos++) {
                if (fem->ja[pos] == i) {
                    fem->matrix[pos] = 1.0;
                } else {
                    fem->matrix[pos] = 0.0;
                }
            }
        }
    }
    
    int boundary_count = 0;
    for (int i = 0; i < fem->num_nodes; i++) {
        boundary_count += fem->boundary_nodes[i];
    }
    printf("Applied boundary conditions to %d nodes\n", boundary_count);
}

// Sparse matrix-vector multiplication (CSR format)
void sparse_matvec(HelmholtzFEM3D* fem, double* x, double* y) {
    for (int i = 0; i < fem->num_nodes; i++) {
        y[i] = 0.0;
        for (int pos = fem->ia[i]; pos < fem->ia[i + 1]; pos++) {
            y[i] += fem->matrix[pos] * x[fem->ja[pos]];
        }
    }
}

// Simple Jacobi solver for sparse system
int solve_jacobi_sparse(HelmholtzFEM3D* fem, double tol, int max_iter) {
    int n = fem->num_nodes;
    double* u_old = malloc(n * sizeof(double));
    double* residual = malloc(n * sizeof(double));
    
    printf("Starting sparse Jacobi solver with tol=%.2e, max_iter=%d\n", tol, max_iter);
    
    // Initialize solution to zero
    memset(fem->solution, 0, n * sizeof(double));
    
    // Find diagonal elements
    double* diag = malloc(n * sizeof(double));
    for (int i = 0; i < n; i++) {
        diag[i] = 1.0;  // Default in case diagonal not found
        for (int pos = fem->ia[i]; pos < fem->ia[i + 1]; pos++) {
            if (fem->ja[pos] == i) {
                diag[i] = fem->matrix[pos];
                break;
            }
        }
    }
    
    for (int iter = 0; iter < max_iter; iter++) {
        // Copy current solution
        memcpy(u_old, fem->solution, n * sizeof(double));
        
        // Jacobi update
        for (int i = 0; i < n; i++) {
            double rhs = fem->rhs[i];
            
            // Subtract off-diagonal terms
            for (int pos = fem->ia[i]; pos < fem->ia[i + 1]; pos++) {
                if (fem->ja[pos] != i) {
                    rhs -= fem->matrix[pos] * u_old[fem->ja[pos]];
                }
            }
            
            fem->solution[i] = rhs / diag[i];
        }
        
        // Check convergence every 10 iterations
        if (iter % 10 == 0) {
            sparse_matvec(fem, fem->solution, residual);
            for (int i = 0; i < n; i++) {
                residual[i] = fem->rhs[i] - residual[i];
            }
            
            double residual_norm = 0.0;
            for (int i = 0; i < n; i++) {
                residual_norm += residual[i] * residual[i];
            }
            residual_norm = sqrt(residual_norm);
            
            printf("Iteration %d: residual norm = %.6e\n", iter, residual_norm);
            
            if (residual_norm < tol) {
                printf("Converged in %d iterations\n", iter);
                free(u_old);
                free(residual);
                free(diag);
                return iter;
            }
        }
    }
    
    printf("Maximum iterations reached without convergence\n");
    free(u_old);
    free(residual);
    free(diag);
    return max_iter;
}

// Initialize FEM solver
HelmholtzFEM3D* helmholtz_fem_init(int nx, int ny, int nz, 
                                   double lx, double ly, double lz,
                                   double kappa) {
    HelmholtzFEM3D* fem = malloc(sizeof(HelmholtzFEM3D));
    
    fem->nx = nx;
    fem->ny = ny;
    fem->nz = nz;
    fem->lx = lx;
    fem->ly = ly;
    fem->lz = lz;
    fem->kappa = kappa;
    
    // Create mesh
    create_structured_tet_mesh(fem);
    
    // Initialize arrays
    fem->rhs = malloc(fem->num_nodes * sizeof(double));
    fem->solution = malloc(fem->num_nodes * sizeof(double));
    
    // Initialize sparse matrix
    init_sparse_matrix(fem);
    
    return fem;
}

// Cleanup FEM solver
void helmholtz_fem_destroy(HelmholtzFEM3D* fem) {
    if (fem) {
        free(fem->nodes);
        free(fem->elements);
        free(fem->matrix);
        free(fem->ia);
        free(fem->ja);
        free(fem->rhs);
        free(fem->solution);
        free(fem->boundary_nodes);
        free(fem);
    }
}

// Test function
int test_helmholtz_fem_3d() {
    printf("=== 3D Helmholtz FEM Solver Test ===\n");
    
    // Problem parameters
    int nx = 16, ny = 16, nz = 16;  // Smaller grid for FEM
    double lx = 1.0, ly = 1.0, lz = 1.0;
    double kappa = 0.1;
    
    printf("Grid: %dx%dx%d, Domain: %.1fx%.1fx%.1f, kappa=%.3f\n", 
           nx, ny, nz, lx, ly, lz, kappa);
    
    // Initialize solver
    HelmholtzFEM3D* fem = helmholtz_fem_init(nx, ny, nz, lx, ly, lz, kappa);
    
    // Assemble system
    assemble_system(fem);
    apply_boundary_conditions(fem);
    
    // Solve
    int iterations = solve_jacobi_sparse(fem, 1e-8, 1000);
    
    // Compute solution norm
    double solution_norm = 0.0;
    for (int i = 0; i < fem->num_nodes; i++) {
        solution_norm += fem->solution[i] * fem->solution[i];
    }
    solution_norm = sqrt(solution_norm);
    
    printf("Solution norm: %.6e\n", solution_norm);
    printf("Iterations used: %d\n", iterations);
    
    // Check matrix properties
    double min_diag = 1e10, max_diag = -1e10;
    int dominant_count = 0;
    
    for (int i = 0; i < fem->num_nodes; i++) {
        if (fem->boundary_nodes[i]) continue;  // Skip boundary nodes
        
        double diag = 0.0, off_diag_sum = 0.0;
        
        for (int pos = fem->ia[i]; pos < fem->ia[i + 1]; pos++) {
            if (fem->ja[pos] == i) {
                diag = fem->matrix[pos];
                min_diag = fmin(min_diag, diag);
                max_diag = fmax(max_diag, diag);
            } else {
                off_diag_sum += fabs(fem->matrix[pos]);
            }
        }
        
        if (diag > off_diag_sum) dominant_count++;
    }
    
    int interior_nodes = fem->num_nodes;
    for (int i = 0; i < fem->num_nodes; i++) {
        interior_nodes -= fem->boundary_nodes[i];
    }
    
    printf("Diagonally dominant interior nodes: %d/%d (%.1f%%)\n", 
           dominant_count, interior_nodes, 100.0 * dominant_count / interior_nodes);
    printf("Diagonal element range: [%.6f, %.6f]\n", min_diag, max_diag);
    
    // Success criteria: reasonable solution norm and convergence
    int success = (solution_norm > 1e-10 && iterations < 1000);
    printf("Test result: %s\n", success ? "PASSED" : "FAILED");
    
    helmholtz_fem_destroy(fem);
    
    return success ? 0 : 1;  // Return 0 for success
}

// Main function for standalone testing
#ifdef HELMHOLTZ_FEM_MAIN
int main(int argc, char** argv) {
    return test_helmholtz_fem_3d();
}
#endif 